
#include "glodef.h"
#include "grid3d.h"
#include "errlog.h"
#include "fe.h"
#include "fe_matr.h"
#include "poisson.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "ns_csplit.h"

/* note: the grid variable used in this file comes from glovar.h */
/*          also flags */
#include "glovar.h"


#define ADJOINT_CURL_FORM            0
#define ADJOINT_REACTION_EXPLICIT    0


#if (ADJOINT_CURL_FORM==1) && (ADJOINT_REACTION_EXPLICIT==1)
#error "only one of ADJOINT_REACTION_EXPLICIT and ADJOINT_CURL_FORM can set to 1"
#endif

/*
 * #define PRINT_ADJOINT_MATRICES
 */


static vector AV_cs, BV_cs; /* global variable, which is local to this file */

typedef struct
{
    unsigned id;
    unsigned eltype[2];
    int ndofs[2];
    vector r;
    vector t;
    scalar d;
    scalar q;
    void *poisson;
    fe_matr_p momentum;
    fe_matr_p adjoint;
} csplit_t;


#define CSPLIT_ID   0xA19E1BBE

static csplit_t *get_cdata(void *data)
{
    csplit_t *p=data;
    if((data==NULL)||(p->id!=CSPLIT_ID))
        error_msg("%s: invalid pointer csplit_t structure passed.\n",__FILE__);
    return p;
}

static csplit_t *_check_input(void *cdata, vector u, scalar p,
    const vector u0, const scalar p0,
    const vector u1, const scalar p1)
{
    csplit_t *cd=get_cdata(cdata);
    if(u==NULL||get_N_vec(u)!=cd->ndofs[0])
        goto _error_;
    if(u0==NULL||get_N_vec(u0)!=cd->ndofs[0])
        goto _error_;
    if(u1==NULL||get_N_vec(u1)!=cd->ndofs[0])
        goto _error_;
    if(p==NULL||get_N_scal(p)!=cd->ndofs[1])
        goto _error_;
    if(p0==NULL||get_N_scal(p0)!=cd->ndofs[1])
        goto _error_;
    if(p1==NULL||get_N_scal(p1)!=cd->ndofs[1])
        goto _error_;
    return cd;
_error_:
    error_msg("Invalid input to TimeStep_cs.\n");
}

/****************************************************************************
 *                                                                          *
 *                                                                          *
 *                                                                          *
 ****************************************************************************
 */

static void get_bdf_ext_coefs(int order, double *bdf, double *ext);

int TimeStep_cs(void *cdata, vector u, scalar p,
    const vector u0, const scalar p0,
    const vector u1, const scalar p1,
    unsigned reassemble)
{
    int ret;
    double tm1, tm, tm_total;

    csplit_t *cd=_check_input(cdata, u, p, u0, p0, u1, p1);
    int i;
    double bdf[5], *ext=bdf+3;
    memset(bdf, 0, sizeof(bdf)); // cleans ext too
    int q_order=flags.time_order;
    if(curr_ti<=q_order)
    {
        q_order=curr_ti;
        reassemble=TRUE;
    }
    get_bdf_ext_coefs(q_order, bdf, ext);

    PrintLog("TimeStep using consistent splitting scheme."
        "(order is %d)\n", q_order);

    /*************************************************************************
     *          STEP 1: The momentum equation
     */

    tm_total=tm=MPI_Wtime();
#define check_point(event)      tm1 = MPI_Wtime(); \
            PrintLog("*** %s: wall time %.3f\n", event, tm1-tm); tm=tm1;

    /*
     * Prepare the right-hand-side of the system
     */
    clean_vec(cd->r);

    /* the source term */
    test_vector_field(grid, cd->eltype[0], TIME, srce_trm, cd->t);
    axpy_vec(cd->r, TStep, cd->t);

    check_point("source term");

    /* the prevoius velocity */
    assert(fabs(bdf[1])>1e-6);
    test_vector(grid, cd->eltype[0], u0, cd->t);
    axpy_vec(cd->r, -bdf[1], cd->t);
    if(fabs(bdf[2])>1e-6)
    {
        test_vector(grid, cd->eltype[0], u1, cd->t);
        axpy_vec(cd->r, -bdf[2], cd->t);
    }

    check_point("previous velocity");

    /* the pressure gradient */
    /*
     * test_grad_e1_in_e2(grid,cd->eltype[1],p0,cd->eltype[0],&cd->t);
     * for(i=0;i<3*cd->ndofs[0];i++)
     * cd->r.X[i]-=TStep*cd->t.X[i];
     */
    assert(fabs(ext[0])>1e-6);
    test_e1_in_div_e2(grid, cd->eltype[1], p0, cd->eltype[0], cd->t);
    axpy_vec(cd->r, TStep*ext[0], cd->t);
    if(fabs(ext[1])>1e-6)
    {
        test_e1_in_div_e2(grid, cd->eltype[1], p1, cd->eltype[0], cd->t);
        axpy_vec(cd->r, TStep*ext[1], cd->t);
    }

    check_point("pressure gradient");

    /*
     * Handle advection
     */
    switch(flags.advection)
    {
        case adv_none: break;
        case adv_explicit:
            clean_vec(AV_cs); // do not modify matrix
            assert(fabs(ext[0])>1e-6);
            test_advection_vector(grid, cd->eltype[0], u0, u0, cd->t);
            axpy_vec(cd->r, -TStep*ext[0], cd->t);
            test_advection_vector_D(grid, cd->eltype[0], u0, u0, cd->t);
            axpy_vec(cd->r, -TStep*ext[0], cd->t);
            if(fabs(ext[1])>1e-6)
            {
                test_advection_vector(grid, cd->eltype[0], u1, u1, cd->t);
                axpy_vec(cd->r, -TStep*ext[1], cd->t);
                test_advection_vector_D(grid, cd->eltype[0], u1, u1, cd->t);
                axpy_vec(cd->r, -TStep*ext[1], cd->t);
            }
            break;
        case adv_semi:
            clean_vec(AV_cs);
            axpy_vec(AV_cs, ext[0], u0);
            if(fabs(ext[1])>1e-6)
            {
                axpy_vec(AV_cs, ext[1], u1);
                /*
                test_advection_vector(grid, cd->eltype[0], u1, u0, cd->t);
                axpy_vec(cd->r, -TStep*ext[1], cd->t);
                test_advection_vector_D(grid, cd->eltype[0], u1, u0, cd->t);
                axpy_vec(cd->r, -TStep*ext[1], cd->t);
                 */
            }
            reassemble=TRUE;
            break;
        default:
            error_msg("%s: advection code %d is either illegal or not implemented\n",
                __func__, flags.advection);
    }
    if(flags.moving_grid&&flags.use_ale)
        error_msg("%s: advection with moving grid and "
        "ALE needs fixing\n", __func__);
    /* TODO : advection with moving grid and ALE needs fixing */

    check_point("advection");

    /*
     * Apply boundary conditions
     */
    apply_bc_dir_vector(grid, cd->eltype[0], TIME, bdry_vel, cd->r);

    {
        /* declared in soah_common.c; Do not expose to everyone */
        extern int on_ival, on_oval;
        if(on_ival|on_oval)
            div_free_bc(grid, cd->eltype[0], on_ival|on_oval, cd->r);
    }
    {
        use_vector(u)
        vector r=cd->r;
        use_vector(r)
        for(i=0; i<cd->ndofs[0]; i++)
        {
            if(fe_dof_bcond(grid, cd->eltype[0], i)==BC_DIR)
            {
                uX[i]=rX[i];
                uY[i]=rY[i];
                uZ[i]=rZ[i];
            }
        }
        done_vector(r)
        done_vector(u)
    }

    check_point("boundary conditions");

    /*
     * Solve the momentum equation
     */
    if(reassemble)
        fe_matr_assemble(grid, 1, &cd->momentum);

    check_point("reassemble matrices");

    PrintLog("Solving the momentum equations.\n");
    /* inital guess for u is u itself (assuming u has the values of u0 */
    if(cd->eltype[0]==fe_bub_vec)
    {
        scalar U=vec2scal(u),
            R=vec2scal(cd->r),
            T=vec2scal(cd->t);
        ret=fe_matr_solve_scalar(cd->momentum, U, R, T, 0.0);
        free_scal(U);
        free_scal(R);
        free_scal(T);
    }
    else
        ret=fe_matr_solve_vector(cd->momentum, u, cd->r, cd->t, 0.0);
    check_point("momentum equation");
    if(!ret)
    {
        warning_msg("sub-step 1 did not converge.\n");
        goto leave_routine;
    }

    /*************************************************************************
     *          STEP 2: The Poisson problem
     */
    /* the velocity time derivative */
    clean_vec(cd->t);
    assert(fabs(bdf[0])>1e-6);
    axpy_vec(cd->t, bdf[0]/TStep, u);
    assert(fabs(bdf[1])>1e-6);
    axpy_vec(cd->t, bdf[1]/TStep, u0);
    if(fabs(bdf[2])>1e-6)
        axpy_vec(cd->t, bdf[2]/TStep, u1);
    test_div_e1_in_e2(grid, cd->eltype[0], cd->t, cd->eltype[1], cd->d);
    clean_scal(cd->q);
    axpy_scal(cd->q, -1.0, cd->d);

    check_point("right-hand-side of Poisson");

    extern scalar phi; // defined in glovar.[hc]
    scalar psi=phi;
    ret=SolvePoisson(cd->poisson, psi, cd->q, cd->d, TRUE, flags.moving_grid);
    check_point("Poisson problem");
    if(!ret)
        goto leave_routine;


    /*************************************************************************
     *          STEP 3: The pressure correction
     */
    /* test psi+p^* and put in cd->q */
    copy_scal(psi, cd->d);
    assert(fabs(ext[0])>1e-6);
    axpy_scal(cd->d, ext[0], p0);
    if(fabs(ext[1])>1e-6)
        axpy_scal(cd->d, ext[1], p1);
    test_scalar(grid, cd->eltype[1], cd->d, cd->q);
    /* test div.u in the pressure element */
    test_div_e1_in_e2(grid, cd->eltype[0], u, cd->eltype[1], cd->d);
    // {
    //     double Coef =  -1.0/GetRe();
    //     for(i=0;i<cd->ndofs[1];i++)
    //         cd->q.F[i] += Coef*cd->d.F[i];
    // } 
    axpy_scal(cd->q, -1.0/GetRe(), cd->d);
    copy_scal(p0, p); /* initial guess */
    /* invert the mass matrix */
    PrintLog("Inverting the pressure mass matrix.\n");
    ret=fe_matr_solve_scalar(massP, p, cd->q, cd->d, 0.0);
    check_point("sub-step 3");
    if(!ret)
    {
        warning_msg("sub-step 3 did not converge.\n");
        goto leave_routine;
    }


#undef check_point
leave_routine:
    PrintLog("******* %s: total wall time = %.3f\n",
        __func__, MPI_Wtime()-tm_total);
    return ret;
}

static void loc_moment(const grid3d_t *grid, int el, unsigned type, double ***S)
{
    static double _M[MAX_LDOFS][MAX_LDOFS];
    static double *M[MAX_LDOFS];
    double mass_coef, stif_coef, advm_coef;
    /* extern vector AV_cs; */
    double **lM, **lS;
    double **lA, **lD;
    double Rho=1.0;
    // double Mu=1.0/GetRe();
    int nn, i, j;
    for(i=0; i<MAX_LDOFS; i++)
        M[i]=_M[i];
    if(S)
        *S=M;
    else
        return;
    double bdf[3]={0.0, 0.0, 0.0};
    int q=flags.time_order;
    if((curr_ti<q)&&(curr_ti>0)) q=curr_ti;
    get_bdf_ext_coefs(q, bdf, NULL);
    // get_coefs(flags.advection, &mass_coef, &stif_coef, &advm_coef);
    mass_coef=Rho*bdf[0];
    stif_coef=1.0/GetRe()*TStep;
    advm_coef=Rho*(flags.advection==adv_semi?TStep:0.0);

    local_mass(grid, el, type, &lM);
    local_stif(grid, el, type, &lS);
    nn=fe_dofs(grid, el, type, NULL);
    for(i=0; i<nn; i++)
        for(j=0; j<nn; j++)
            (*S)[i][j]=mass_coef*lM[i][j]+stif_coef*lS[i][j];
    if(flags.advection==adv_semi)
    {
        local_advm(grid, el, type, AV_cs, &lA);
        local_advm_D(grid, el, type, AV_cs, &lD);
        for(i=0; i<nn; i++)
            for(j=0; j<nn; j++)
                (*S)[i][j]+=advm_coef*(lA[i][j]+lD[i][j]);
    }
}

static void loc_adjoint(const grid3d_t *grid, int el, unsigned type, double ***S);

void *InitCSplit(unsigned elemU, unsigned elemP)
{
    csplit_t *cd;
    int symmetric;
    double tm=MPI_Wtime();

    InitFE(1, &elemU, grid);
    PrintLog("Initialize consistent splitting scheme (%s) begin.\n", __func__);

    cd=malloc(sizeof(csplit_t));
    if(cd==NULL)
        error_msg("ns_csplit.c: memory allocation error.\n");
    assert(grid!=NULL);
    cd->id=CSPLIT_ID;
    cd->eltype[0]=elemU;
    cd->eltype[1]=elemP;
    if(elemU==fe_bub_vec)
        cd->ndofs[0]=fe_ndof(grid, fe_bub);
    else
        cd->ndofs[0]=fe_ndof(grid, elemU);
    cd->ndofs[1]=fe_ndof(grid, elemP);
    AV_cs=alloc_vec(cd->ndofs[0]);
    BV_cs=alloc_vec(cd->ndofs[0]);
    cd->r=alloc_vec(cd->ndofs[0]);
    cd->t=alloc_vec(cd->ndofs[0]);
    cd->q=alloc_scal(cd->ndofs[1]);
    cd->d=alloc_scal(cd->ndofs[1]);
    symmetric=flags.advection!=adv_semi;
    cd->momentum=fe_matr_alloc(elemU, loc_moment, eo_bc_dir, grid, symmetric,
#if defined(USE_CSC)
        symmetric?"sp":"uu"
#else
        NULL
#endif
        );
    // fe_matr_structure(grid, 1, &cd->momentum);
    fe_matr_assemble(grid, 1, &cd->momentum);

    if(flags.need_adjoint)
    {
        unsigned elemA=elemU;
#if (ADJOINT_REACTION_EXPLICIT==1)
        PrintLog("\n => ADJOINT_REACTION_EXPLICIT is 1\n");
#elif (ADJOINT_CURL_FORM==1)
        PrintLog("\n => ADJOINT_CURL_FORM is 1\n");
        if(flags.adj_advection==adv_semi)
        {
            if(elemU!=fe_bub&&elemU!=fe_bub_vec)
                error_msg("Adjoint N-S solver in curl form with semi-implicit advection "
                "is not available for %s element\n",
                fe_long_names[elemU]);
            else
                elemA=fe_bub_vec;
        }
#else
        if(flags.adj_advection==adv_semi)
        {
            if(elemU!=fe_bub&&elemU!=fe_bub_vec)
                error_msg("Adjoint N-S solver in standard form with semi-implicit advection "
                "is not available for %s element\n",
                fe_long_names[elemU]);
            else
                elemA=fe_bub_vec;
        }
#endif
        // PRINTF("elemA = %d\n", elemA);
        unsigned symm=flags.adj_advection!=adv_semi;
        cd->adjoint=fe_matr_alloc(elemA, loc_adjoint,
            eo_bc_dir, grid, symm,
#if defined(USE_CSC) 
            symm?"sp":"uu"
#else
            NULL
#endif
            );
        // fe_matr_structure(grid, 1, &cd->adjoint);
        fe_matr_assemble(grid, 1, &cd->adjoint);
    }
    else
        cd->adjoint=NULL;

    // PRINTF("get_el_type(cd->momentum) = %d\n", get_el_type(cd->momentum));
    // PRINTF("get_el_type(cd->adjoint) = %d\n", get_el_type(cd->adjoint));

    PrintLog("%s: total (excluding InitPoisson) CPU time  %-6.3fsec.\n", __func__,
        (MPI_Wtime()-tm));

    cd->poisson=InitPoisson(grid, elemP);
    return cd;
}

void DoneCSplit(void *cdata)
{
    csplit_t *cd=get_cdata(cdata);
    DonePoisson(cd->poisson);
    free_vec(AV_cs);
    free_vec(BV_cs);
    free_vec(cd->r);
    free_vec(cd->t);
    free_scal(cd->q);
    free_scal(cd->d);
    fe_matr_free(cd->momentum);
    fe_matr_free(cd->adjoint);
    free(cd);

    PrintLog("Consistent splitting scheme done.\n");
}

/****************************************************************************/

static void loc_adjoint(const grid3d_t *grid, int el, unsigned type, double ***S)
{
    static double _M[MAX_LDOFS][MAX_LDOFS];
    static double *M[MAX_LDOFS];
    double mass_coef, stif_coef, advm_coef;
    /* extern vector AV_cs; */
    double **lM, **lS;

    double Rho=1.0;
    // double Mu=1.0/GetRe();
    int nn, i, j;
    memset(_M, 0, sizeof(_M));
    for(i=0; i<MAX_LDOFS; i++)
        M[i]=_M[i];
    if(S)
        *S=M;
    else
        return;
    double bdf[3]={0.0, 0.0, 0.0};
    int q=flags.time_order;
    if(((total_ti-curr_ti)<q)&&(curr_ti<total_ti)) q=total_ti-curr_ti;
    get_bdf_ext_coefs(q, bdf, NULL);
    // get_coefs(flags.adj_advection, &mass_coef, &stif_coef, &advm_coef);
    mass_coef=Rho*bdf[0];
    stif_coef=1.0/GetRe()*TStep;
    advm_coef=Rho*(flags.adj_advection==adv_semi?TStep:0.0);

    local_mass(grid, el, type, &lM);
    local_stif(grid, el, type, &lS);
    nn=fe_dofs(grid, el, type, NULL);
    for(i=0; i<nn; i++)
        for(j=0; j<nn; j++)
            (*S)[i][j]=mass_coef*lM[i][j]+stif_coef*lS[i][j];
    if(flags.adj_advection==adv_semi)
#if (ADJOINT_REACTION_EXPLICIT==1)
    {
        double **lG, **lD;
        local_advm(grid, el, type, AV_cs, &lG);
        local_advm_D(grid, el, type, AV_cs, &lD);
        for(i=0; i<nn; i++)
            for(j=0; j<nn; j++)
                (*S)[i][j]+=advm_coef*(-lG[i][j]-lD[i][j]);
    }
#elif (ADJOINT_CURL_FORM==1)
        {
            double **la, **lu, **lR;
            local_advm(grid, el, type, AV_cs, &la);
            local_advm(grid, el, type, BV_cs, &lu);
            local_adjoit_curl(grid, el, type, AV_cs, &lR);
            for(i=0; i<nn; i++)
                for(j=0; j<nn; j++)
                    (*S)[i][j]+=advm_coef*(-la[i][j]+lu[i][j]+lR[i][j]);
        }
#else
        {
            double **lG, **lD, **lB;
            local_advm(grid, el, type, AV_cs, &lG);
            local_advm_D(grid, el, type, AV_cs, &lD); // stabilization of advection ??
            local_adjoit_B(grid, el, type, BV_cs, &lB);
            // local_adjoit_C(grid, el, type, BV_cs, &lC);
            for(i=0; i<nn; i++)
                for(j=0; j<nn; j++)
                    (*S)[i][j]+=advm_coef*(-lG[i][j]-lD[i][j]+lB[i][j]);
            //        (*S)[i][j] +=  Rho*advm_coef*(-lG[i][j]-lD[i][j]+lB[i][j]+lC[i][j]);
        }
#endif
}

#if defined(PRINT_ADJOINT_MATRICES)

static void loc_G(const grid3d_t *grid, int el, unsigned type, double ***S)
{
    local_advm(grid, el, type, AV_cs, S);
}

static void loc_D(const grid3d_t *grid, int el, unsigned type, double ***S)
{
    local_advm_D(grid, el, type, AV_cs, S);
}

static void loc_B(const grid3d_t *grid, int el, unsigned type, double ***S)
{
    local_adjoit_B(grid, el, type, BV_cs, S);
}

static void loc_C(const grid3d_t *grid, int el, unsigned type, double ***S)
{
    local_adjoit_C(grid, el, type, BV_cs, S);
}
#endif

/****************************************************************************/

int TimeStep_adjoint_cs(void *cdata, vector v, scalar q,
    const vector v0, const scalar q0,
    const vector v1, const scalar q1,
    vector adv_u, vector u,
    vector f, vector vbc,
    int test_f)
/*
 *  If we are solving for V^m, then set
 *      v0 = V^{m+1},
 *      adv_u = U^m + ndisp/TStep,
 *      u = U^m,
 *  where ndisp is a returned by get_ndisp in alegrid.c, i.e.
 *      ndisp = x^m-x^{m+1} = grid->nodes - alegrid->old_nodes,
 *
 *  NOTE: ndisp here is -ndisp of the state problem, since time is moving backwards.
 */
{
    int ret;
    double tm1, tm, tm_total;

    csplit_t *cd=_check_input(cdata, v, q, v0, q0, v1, q1);
    int i, reassemble, q_order;
    if(flags.moving_grid) reassemble=TRUE;
    else reassemble=FALSE;

    double bdf[5], *ext=bdf+3;
    memset(bdf, 0, sizeof(bdf)); // cleans ext too
    q_order=flags.time_order;
    if(q_order>=total_ti-curr_ti)
    {
        q_order=total_ti-curr_ti;
        reassemble=TRUE;
    }
    get_bdf_ext_coefs(q_order, bdf, ext);

    PrintLog("Adjoint TimeStep using consistent splitting scheme."
        "(order is %d)\n", q_order);

    /*************************************************************************
     *          STEP 1: The momentum equation
     */

    tm_total=tm=MPI_Wtime();
#define check_point(event)      tm1 = MPI_Wtime(); \
            PrintLog("*** %s: wall time %.3f\n", event, tm1-tm); tm=tm1;

    /*
     * Prepare the right-hand-side of the system
     */

    /* the source term */
    if(f)
    {
        if(test_f)
        {
            test_vector(grid, cd->eltype[0], f, cd->t);
            // for(i=0;i<3*cd->ndofs[0];i++) cd->r.X[i]=TStep*cd->t.X[i];
            axpy_vec(cd->r, TStep, cd->t);
        }
        else
            copy_vec(f, cd->r);
    }
    else
        clean_vec(cd->r);
    // for(i=0;i<3*cd->ndofs[0];i++) cd->r.X[i] *= TStep;
    scale_vec(cd->r, TStep);

    check_point("source term");

    /* the prevoius adjoint velocity */
    assert(fabs(bdf[1])>1e-6);
    test_vector(grid, cd->eltype[0], v0, cd->t);
    axpy_vec(cd->r, -bdf[1], cd->t);
    if(fabs(bdf[2])>1e-6)
    {
        test_vector(grid, cd->eltype[0], v1, cd->t);
        axpy_vec(cd->r, -bdf[2], cd->t);
    }

    check_point("previous velocity");

    /* the adjoint pressure gradient */
    /*
     * test_grad_e1_in_e2(grid,cd->eltype[1],p0,cd->eltype[0],&cd->t);
     * for(i=0;i<3*cd->ndofs[0];i++)
     * cd->r.X[i]-=TStep*cd->t.X[i];
     */
    assert(fabs(ext[0])>1e-6);
    test_e1_in_div_e2(grid, cd->eltype[1], q0, cd->eltype[0], cd->t);
    // for(i=0;i<3*cd->ndofs[0];i++) cd->r.X[i]+=TStep*cd->t.X[i];
    axpy_vec(cd->r, TStep*ext[0], cd->t);
    if(fabs(ext[1])>1e-6)
    {
        test_e1_in_div_e2(grid, cd->eltype[1], q1, cd->eltype[0], cd->t);
        axpy_vec(cd->r, TStep*ext[1], cd->t);
    }

    check_point("pressure gradient");

    /*
     * Handle advection
     */
    if(adv_u==NULL)
        adv_u=u;
    switch(flags.adj_advection)
    {
        case adv_none: /* do nothing */
            break;
        case adv_explicit: /* explicit advection -- modify cd->r */
            clean_vec(AV_cs); /* keep matrix as is */
            clean_vec(BV_cs);
            assert(fabs(ext[0])>1e-6);
            test_advection_vector(grid, cd->eltype[0], adv_u, v0, cd->t);
            // for(i=0;i<3*cd->ndofs[0];i++) cd->r.X[i]+=TStep*cd->t.X[i];
            axpy_vec(cd->r, TStep*ext[0], cd->t);
            if(fabs(ext[1])>1e-6)
            {
                test_advection_vector(grid, cd->eltype[0], adv_u, v1, cd->t);
                axpy_vec(cd->r, TStep*ext[1], cd->t);
            }

            assert(cd->eltype[0]==fe_bub||cd->eltype[0]==fe_bub_vec);
            assert(fabs(ext[0])>1e-6);
            test_adjoint_B_vector(grid, fe_bub_vec, u, v0, cd->t);
            // for(i=0;i<3*cd->ndofs[0];i++) cd->r.X[i]-=TStep*cd->t.X[i];
            axpy_vec(cd->r, -TStep*ext[0], cd->t);
            if(fabs(ext[1])>1e-6)
            {
                test_adjoint_B_vector(grid, fe_bub_vec, u, v1, cd->t);
                axpy_vec(cd->r, -TStep*ext[1], cd->t);
            }
            break;
        case adv_semi: /* semi-implicit advection -- modify matrix */
            copy_vec(adv_u, AV_cs);
            copy_vec(u, BV_cs);
#if (ADJOINT_REACTION_EXPLICIT==1)
            assert(cd->eltype[0]==fe_bub||cd->eltype[0]==fe_bub_vec);
            assert(fabs(ext[0])>1e-6);
            test_adjoint_B_vector(grid, fe_bub_vec, u, v0, cd->t);
            // for(i=0;i<3*cd->ndofs[0];i++) cd->r.X[i]-=TStep*cd->t.X[i];
            axpy_vec(cd->r, -TStep*ext[0], cd->t);
            if(fabs(ext[1])>1e-6)
            {
                test_adjoint_B_vector(grid, fe_bub_vec, u, v1, cd->t);
                axpy_vec(cd->r, -TStep*ext[1], cd->t);
            }
            clean_vec(BV_cs);
#endif
            reassemble=TRUE; /* must reassemble matrix */
            break;
        default:
            error_msg("%s: adjoint advection code %d is either illegal or "
                "not implemented\n", __func__, flags.adj_advection);
    }

    check_point("advection");

    /*
     * Apply boundary conditions
     */
    {
        vector r=cd->r;
        use_vector(r)
        use_vector(v)
        if(vbc==NULL)
        {
            for(i=0; i<cd->ndofs[0]; i++)
            {
                if(fe_dof_bcond(grid, cd->eltype[0], i)==BC_DIR)
                {
                    vX[i]=rX[i]=0.0;
                    vY[i]=rY[i]=0.0;
                    vZ[i]=rZ[i]=0.0;
                }
            }
        }
        else
        {
            use_vector(vbc)
            for(i=0; i<cd->ndofs[0]; i++)
            {
                if(fe_dof_bcond(grid, cd->eltype[0], i)==BC_DIR)
                {
                    vX[i]=rX[i]=vbcX[i];
                    vY[i]=rY[i]=vbcY[i];
                    vZ[i]=rZ[i]=vbcZ[i];
                }
            }
            done_vector(vbc)
        }
        done_vector(v)
        done_vector(r)
    }

    check_point("boundary conditions");

#if defined(PRINT_ADJOINT_MATRICES)
    if(1)
    {
        const int num_mats=5;
        fe_matr_t*mats[num_mats];
        char fname[1024];
        int digits=lround(floor(log10(total_ti)+1));

        mats[0]=fe_matr_alloc(fe_bub_vec, loc_G, eo_add_all, grid, FALSE);
        mats[1]=fe_matr_alloc(fe_bub_vec, loc_D, eo_add_all, grid, FALSE);
        mats[2]=fe_matr_alloc(fe_bub_vec, loc_B, eo_add_all, grid, FALSE);
        mats[3]=fe_matr_alloc(fe_bub_vec, loc_C, eo_add_all, grid, FALSE);
        mats[4]=fe_matr_alloc(fe_bub_vec, local_mass, eo_add_all, grid, FALSE);

        fe_matr_assemble(grid, num_mats, mats);

        sprintf(fname, "mats/adv_G_s_%0*d.txt", digits, curr_ti);
        PrintSMatrix(mats[0]->matr, fname);
        sprintf(fname, "mats/adv_G_b_%0*d.txt", digits, curr_ti);
        PrintSMatrix(mats[0]->matr_bc, fname);

        sprintf(fname, "mats/adv_D_s_%0*d.txt", digits, curr_ti);
        PrintSMatrix(mats[1]->matr, fname);
        sprintf(fname, "mats/adv_D_b_%0*d.txt", digits, curr_ti);
        PrintSMatrix(mats[1]->matr_bc, fname);

        sprintf(fname, "mats/adv_B_s_%0*d.txt", digits, curr_ti);
        PrintSMatrix(mats[2]->matr, fname);
        sprintf(fname, "mats/adv_B_b_%0*d.txt", digits, curr_ti);
        PrintSMatrix(mats[2]->matr_bc, fname);

        sprintf(fname, "mats/adv_C_s_%0*d.txt", digits, curr_ti);
        PrintSMatrix(mats[3]->matr, fname);
        sprintf(fname, "mats/adv_C_b_%0*d.txt", digits, curr_ti);
        PrintSMatrix(mats[3]->matr_bc, fname);

        sprintf(fname, "mats/mass_s_%0*d.txt", digits, curr_ti);
        PrintSMatrix(mats[4]->matr, fname);
        sprintf(fname, "mats/mass_b_%0*d.txt", digits, curr_ti);
        PrintSMatrix(mats[4]->matr_bc, fname);

        for(int i=0; i<num_mats; i++)
            fe_matr_free(mats[i]);
    }
#endif


    /*
     * Solve the momentum equation
     */
    if(reassemble)
        fe_matr_assemble(grid, 1, &cd->adjoint);

    check_point("reassemble matrices");

    PrintLog("Solving the adjoint momentum equations.\n");
    if(get_el_type(cd->adjoint)==fe_bub_vec)
    {
        scalar V=vec2scal(v),
            R=vec2scal(cd->r),
            T=vec2scal(cd->t);
        ret=fe_matr_solve_scalar(cd->adjoint, V, R, T, 0.0);
        free_scal(V);
        free_scal(R);
        free_scal(T);
    }
    else
        ret=fe_matr_solve_vector(cd->adjoint, v, cd->r, cd->t, 0.0);
    check_point("momentum equation");
    if(!ret)
    {
        warning_msg("sub-step 1 did not converge.\n");
        goto leave_routine;
    }



    /*************************************************************************
     *          STEP 2: The Poisson problem
     */
    /* the adjoint velocity time derivative */
    clean_vec(cd->t);
    assert(fabs(bdf[0])>1e-6);
    axpy_vec(cd->t, bdf[0]/TStep, v);
    assert(fabs(bdf[1])>1e-6);
    axpy_vec(cd->t, bdf[1]/TStep, v0);
    if(fabs(bdf[2])>1e-6)
        axpy_vec(cd->t, bdf[2]/TStep, v1);
    test_div_e1_in_e2(grid, cd->eltype[0], cd->t, cd->eltype[1], cd->d);
    // for(i=0;i<cd->ndofs[1];i++) cd->q.F[i] = -cd->d.F[i];
    clean_scal(cd->q);
    axpy_scal(cd->q, -1.0, cd->d);

    check_point("right-hand-side of Poisson");

    extern scalar qhi; // defined in glovar.[hc]
    scalar psi=qhi;
    ret=SolvePoisson(cd->poisson, psi, cd->q, cd->d, TRUE, flags.moving_grid);
    check_point("Poisson problem");
    if(!ret)
        goto leave_routine;

    /*************************************************************************
     *          STEP 3: The pressure correction
     */
    /* test psi+q^* and put in cd->q */
    copy_scal(psi, cd->d);
    assert(fabs(ext[0])>1e-6);
    axpy_scal(cd->d, ext[0], q0);
    if(fabs(ext[1])>1e-6)
        axpy_scal(cd->d, ext[1], q1);
    test_scalar(grid, cd->eltype[1], cd->d, cd->q);
    /* test div.v in the pressure element */
    test_div_e1_in_e2(grid, cd->eltype[0], v, cd->eltype[1], cd->d);
    // {
    //     double Coef =  -1.0/GetRe();
    //     for(i=0;i<cd->ndofs[1];i++)
    //         cd->q.F[i] += Coef*cd->d.F[i];
    // }
    axpy_scal(cd->q, -1.0/GetRe(), cd->d);
    copy_scal(q0, q); /* initial guess */
    /* invert the mass matrix */
    PrintLog("Inverting the pressure mass matrix.\n");
    ret=fe_matr_solve_scalar(massP, q, cd->q, cd->d, 0.0);
    check_point("sub-step 3");
    if(!ret)
    {
        warning_msg("sub-step 3 did not converge.\n");
        goto leave_routine;
    }

#undef check_point
leave_routine:
    PrintLog("******* %s: total wall time = %.3f\n",
        __func__, MPI_Wtime()-tm_total);
    return ret;
}


#include "fe_bubble.h"
#include "vecs.h"

void InitialPressure_cs(void* cdata, vector u0, vector ut0_bc,
    scalar p)
{
    csplit_t *cd=get_cdata(cdata);

    double tm=MPI_Wtime();

    PrintLog("Solving for initial pressure.    ");

    if(((cd->eltype[0]!=fe_bub)&&(cd->eltype[0]!=fe_bub_vec))||(cd->eltype[1]!=fe_p1))
        error_msg("%s is only available in P1-bubble element.\n", __func__);

    /* the source term */
    interpolate_vector_lagr_bub(grid, 0.0, srce_trm, cd->t);
    test_bub_in_grad_p1(grid, cd->t, p);

    copy_scal(p, cd->d);

    if(flags.advection!=adv_none)
    {
        /* as a test: this should be equivalent to the two lines that follow */
#if FALSE
        test_advection_vector(grid, fe_bub, u0, u0, BV_cs);
        fe_matr_solve_vector(massU, cd->r, BV_cs, AV_cs, 0.0);
        test_bub_in_grad_p1(grid, cd->r, p);
#endif
        /* the advection term */
        test_bub_grad_bub_in_grad_p1(grid, u0, u0, p);
        // for(i=0;i<p->N;i++) cd->d.F[i] -= p->F[i];
        axpy_scal(cd->d, -1.0, p);
    }

    /* the diffusion term */
    surf_test_n_cross_curl_bub_in_grad_p1(grid, u0, p);
    // for(i=0;i<p->N;i++) cd->d.F[i] -= p->F[i]/Re;
    axpy_scal(cd->d, -1.0/GetRe(), p);

    /* the time derivative */
    surf_test_bub_dot_n_in_p1(grid, ut0_bc, p);
    // for(i=0;i<p->N;i++) cd->d.F[i] -= p->F[i];
    axpy_scal(cd->d, -1.0, p);



    PrintLog("CPU time for assembly of right-hand-side: %-6.3fsec.\n",
        (MPI_Wtime()-tm));

    clean_scal(p);
    SolvePoisson(cd->poisson, p, cd->d, cd->q, TRUE, TRUE);

    PrintLog("    CPU time for inital pressure: %-6.3fsec.\n",
        (MPI_Wtime()-tm));

}

void FinalAdjointPressure_cs(void* cdata, vector vT, vector uT,
    vector vtT_bc, vector f, scalar f_tested,
    scalar r)
{
    csplit_t *cd=get_cdata(cdata);

    double tm=MPI_Wtime();

    PrintLog("Solving for final adjoint pressure.    ");

    if(((cd->eltype[0]!=fe_bub)&&(cd->eltype[0]!=fe_bub_vec))||(cd->eltype[1]!=fe_p1))
        error_msg("%s is only available in P1-bubble element.\n", __func__);

    /* the source term */
    if(f_tested)
        copy_scal(f_tested, r);
    else if(f)
        test_bub_in_grad_p1(grid, f, r);
    else
        clean_scal(r);

    copy_scal(r, cd->d);

    if(flags.adj_advection!=adv_none)
    {
        /* the adjoint advection terms */
        test_bub_grad_bub_in_grad_p1(grid, uT, vT, r);
        // for(i=0;i<r->N;i++) cd->d.F[i] += r->F[i];
        axpy_scal(cd->d, 1.0, r);
        test_bub_grad_bub_in_grad_p1(grid, vT, uT, r);
        // for(i=0;i<r->N;i++) cd->d.F[i] -= r->F[i];
        axpy_scal(cd->d, -1.0, r);
    }

    /* the diffusion term */
    surf_test_n_cross_curl_bub_in_grad_p1(grid, vT, r);
    // for(i=0;i<r->N;i++) cd->d.F[i] -= r->F[i]/Re;
    axpy_scal(cd->d, -1.0/GetRe(), r);

    /* the time derivative */
    surf_test_bub_dot_n_in_p1(grid, vtT_bc, r);
    // for(i=0;i<r->N;i++) cd->d.F[i] += r->F[i];
    axpy_scal(cd->d, 1.0, r);

    PrintLog("CPU time for assembly of right-hand-side: %-6.3fsec.\n",
        (MPI_Wtime()-tm));

    clean_scal(r);
    SolvePoisson(cd->poisson, r, cd->d, cd->q, TRUE, TRUE);

    PrintLog("    CPU time for inital pressure: %-6.3fsec.\n",
        (MPI_Wtime()-tm));

}

static void get_bdf_ext_coefs(int order, double *bdf, double *ext)
{
    if(order>2)
        error_msg("%s: time integration of "
        "order %d isn't available\n", __FILE__, order);
    switch(order)
    {
        case 1:
            if(bdf)
            {
                bdf[0]=1.0;
                bdf[1]= -1.0;
            }
            if(ext)
            {
                ext[0]=1.0;
            }
            return;
        case 2:
            if(bdf)
            {
                bdf[0]=1.5;
                bdf[1]= -2.0;
                bdf[2]=0.5;
            }
            if(ext)
            {
                ext[0]=2.0;
                ext[1]= -1.0;
            }
            return;
        default:
            error_msg("Invalid order");
    }
}


